##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--mut_rate_gene_file"), type = "character") ,
    make_option(c("--smg_file"), type = "character") ,
    make_option(c("--mol_type"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    mol_type <- "MSI"
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    mut_rate_gene_file <- paste(work_dir,"/images/mutRate/MutRate.molType.tsv",sep="")
    smg_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutsig_check/smg.list"
	images_path <- paste(work_dir,"/images/mutRatePlot",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

mut_rate_gene_file <- opt$mut_rate_gene_file
smg_file <- opt$smg_file
images_path <- opt$images_path
mol_type <- opt$mol_type

dir.create(images_path , recursive = T)

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")
col_im_gc <- col[c(1,4)]
col_igc_dgc <- col[c(2,3)]

###########################################################################################

dat_mutRateGene <- data.frame(fread( mut_rate_gene_file ))
smg <- data.frame(fread(smg_file))$Gene_Symbol

###########################################################################################

dat_mutRateGene <- subset(dat_mutRateGene , Hugo_Symbol %in% smg)

dat_plot <- rbind( dat_mutRateGene )
dat_plot$Class <- factor( dat_plot$Class , levels = c("IM" , "GC" , "IGC" , "DGC") , order = T )
dat_plot$value_text <- paste0( round(dat_plot$MutRate , 2) * 100 , "%") 

###########################################################################################
if(mol_type=="MSI"){
    dat_plot <- subset( dat_plot , Molecular.subtype == "MSI/POLE" )
}else{
    dat_plot <- subset( dat_plot , Molecular.subtype == mol_type )
}

im_num <- unique(subset(dat_plot , Class=="IM")$SampleNum)
gc_num <- unique(subset(dat_plot , Class=="GC")$SampleNum)
igc_num <- unique(subset(dat_plot , Class=="IGC")$SampleNum)
dgc_num <- unique(subset(dat_plot , Class=="DGC")$SampleNum)

###########################################################################################
## 计算P值
dat_plot$p.value = ""
dat_plot$p_text = ""

for(geneN in unique(dat_plot$Hugo_Symbol)){

    print(geneN)

    tmp_1 <- subset( dat_plot , Hugo_Symbol == geneN & Class %in% c("GC") )
    tmp_2 <- subset( dat_plot , Hugo_Symbol == geneN & Class %in% c("IM") )

    if(nrow(tmp_1)==0){
        tmp_1 <- tmp_2
        tmp_1$MutNum <- 0
        tmp_1$MutRate <- 0
        tmp_1$value_text <- ""
    }

    if(nrow(tmp_2)==0){
        tmp_2 <- tmp_1
        tmp_2$MutNum <- 0
        tmp_2$MutRate <- 0
        tmp_2$value_text <- ""
    }

    tmp <- rbind( tmp_1 , tmp_2 )

    tmp_fisher <- matrix(c(tmp$MutNum , tmp$SampleNum - tmp$MutNum) , ncol = 2)
    p <- fisher.test(tmp_fisher)$p.value

    dat_plot[dat_plot$Hugo_Symbol == geneN & dat_plot$Class %in% c("GC" , "IM"),"p.value"] <- p
    p_text <- format(as.numeric(p) , scientific = TRUE , digits = 3)
    dat_plot[dat_plot$Hugo_Symbol == geneN & dat_plot$Class %in% c("GC" , "IM"),"p_text"] <- p_text
}

for(geneN in unique(dat_plot$Hugo_Symbol)){

    print(geneN)

    tmp_1 <- subset( dat_plot , Hugo_Symbol == geneN & Class %in% c("IGC") )
    tmp_2 <- subset( dat_plot , Hugo_Symbol == geneN & Class %in% c("DGC") )

    if(nrow(tmp_1)==0){
        tmp_1 <- tmp_2
        tmp_1$MutNum <- 0
        tmp_1$MutRate <- 0
        tmp_1$value_text <- ""
    }

    if(nrow(tmp_2)==0){
        tmp_2 <- tmp_1
        tmp_2$MutNum <- 0
        tmp_2$MutRate <- 0
        tmp_2$value_text <- ""
    }

    tmp <- rbind( tmp_1 , tmp_2 )

    tmp_fisher <- matrix(c(tmp$MutNum , tmp$SampleNum - tmp$MutNum) , ncol = 2)
    p <- fisher.test(tmp_fisher)$p.value

    dat_plot[dat_plot$Hugo_Symbol == geneN & dat_plot$Class %in% c("IGC" , "DGC"),"p.value"] <- p
    p_text <- format(as.numeric(p) , scientific = TRUE , digits = 3)
    dat_plot[dat_plot$Hugo_Symbol == geneN & dat_plot$Class %in% c("IGC" , "DGC"),"p_text"] <- p_text
}

gene_order <- c("TP53" , "ARID1A" , "CDH1" , "APC" , "SMAD4" , "DNAH3" , "MUC6" , "PIK3CA" , "CTNNB1" , 
    "RHOA" , "ERBB2" , "CFTR" , "KRAS" , "MAP2K7" , "ARID2" , "RNF43" , "TGFBR2" ,
    "BMP6" , "FBXW7" , "CDKN2A" , "MTRR" , "MICAL2")

dat_plot$Hugo_Symbol <- factor( dat_plot$Hugo_Symbol , 
    levels = gene_order , order = T)

dat_plot <- subset( dat_plot )

###########################################################################################

plotRate <- function( result_use = result_use , col_use = col_use, images_name = images_name , title = title , width = width , height = height ){
    ###########################################################################################
    ## 柱状图展示
    ## p值
    result_use$p_text=ifelse(result_use$p>=0.05,"","*")
    result_use$p_text=ifelse(result_use$p<0.05 & result_use$p>0.01,"*",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.01 & result_use$p>0.001,"**",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.001 ,"***",result_use$p_text)

    ##########################################################################################

    result2 <- result_use
    result2$value_percent=paste(round(result2$MutRate * 100),"%",sep="")
    result2$variable <- result2$Class

    p <- ggplot(result2,mapping = aes(variable,MutRate,fill=variable))+
    geom_bar(stat='identity',position='stack',color = 'black') + 
    facet_grid(.~gene)+
    theme_bw() +
    labs(title = title) +
    scale_fill_manual(values=col) +
    theme(strip.text.x = element_text(size = 11))+
    theme(panel.grid=element_blank())+labs(x = 'Genes',y = 'Mutation Rate') +
    theme(axis.title =element_text(size = 15),axis.text =element_text(size = 14, color = 'black'))+
    theme(axis.ticks.x = element_blank(),axis.text.x = element_blank())+
    geom_text(aes(label=value_percent, y=MutRate+0.003), position=position_dodge(0.9), vjust=0)+
    geom_text(aes(label=p_text, y=0.8,x=1.5),size=7)+
    xlab("") +
    scale_fill_manual(values=col_use) +
    ylim(0,0.8) +
    theme(
      title =element_text(size=4, face='bold'),
      legend.title = element_blank(),
      legend.text = element_text(size = 12),
      legend.key.width = unit(1, "cm"),
      legend.key.height = unit(1, "cm"),
      plot.title = element_text(size = 30, face = "bold")
    )

    ggsave( images_name , p , width = width , height = height )

}

result_use <- subset( dat_plot , Class %in% c("IGC" , "DGC") )
result_use$p <- as.numeric(result_use$p.value)
result_use$gene <- result_use$Hugo_Symbol
#title <- paste0( mol_type , " (" , "IGC:" , igc_num , " , " , "DGC:" , dgc_num , ")" )
title <- mol_type
width <- 18
height <- 6
col_use <- col_igc_dgc
images_name <- paste0(images_path , "/MutRate.compare.IGC_DGC.",mol_type,".pdf")
plotRate( result_use = result_use , col_use = col_use, images_name = images_name , title = title , width = width , height = height )

